#!/usr/bin/python
'''
Finds the overlap between two GFF files.
'''
import sys
import os
import numpy as np
import overlap

# parameters.
prec_ov = 0.0000000001
just_mark = False
hits = False
misses = False
for i in range(len(sys.argv)):
	if sys.argv[i] == "-t":
		target_file = sys.argv[i+1]
	elif sys.argv[i] == "-q":
		query_file = sys.argv[i+1]
	elif sys.argv[i] == "-hits":
		hits = True
		target_out_file = sys.argv[i+1]
		query_out_file = sys.argv[i+2]
	elif sys.argv[i] == "-misses":
		misses = True
		target_out_file = sys.argv[i+1]
		query_out_file = sys.argv[i+2]
	elif sys.argv[i] == "-p":
		prec_ov = float(sys.argv[i+1])
	elif sys.argv[i] == "-mark":
		just_mark = sys.argv[i+1]

#debug_sz = 1000000

################################
# read target file into memory.
print "reading target."
fin = open(target_file, "rb")
#lines = fin.readlines(debug_sz)
lines = fin.readlines()
fin.close()

# count number of non-comment lines.
target_sz = 0
for line in lines:
	if line[0] != "#": target_sz += 1

# parse target file.
print "parsing targets.", target_sz
j = 0
#targets = np.zeros(len(lines), dtype=overlap.entry_dt)
targets = np.zeros(target_sz, dtype=overlap.entry_dt)
target_lines = []
refs = {}
idx = 0
for i in range(len(lines)):
	# skips and token.
	if lines[i][0] == "#": continue
	tmp = lines[i].strip().split("\t")
	if len(tmp) != 9: continue
	target_lines.append(tmp)
	
	# note refs.
	if tmp[0] not in refs:
		refs[tmp[0]] = idx
		idx += 1
	
	# add to array.
	targets[j]['type'] = True
	targets[j]['ref'] = refs[tmp[0]]
	targets[j]['start'] = int(tmp[3])
	targets[j]['stop'] = int(tmp[4])
	targets[j]['info'] = j
	j += 1
	
	#if i >= debug_sz: break


# read in node file.
print "reading query file."
fin = open(query_file, "rb")
#lines = fin.readlines(debug_sz)
lines = fin.readlines()
fin.close()

# count number of non-comment lines.
query_sz = 0
for line in lines:
	if line[0] != "#": query_sz += 1


# parse queries
print "parsing queries.", query_sz
j = 0
queries = np.zeros(query_sz, dtype=overlap.entry_dt)
query_lines = []
for i in range(len(lines)):
	# skips and token.
	if lines[i][0] == "#": continue
	tmp = lines[i].strip().split("\t")
	query_lines.append(tmp)
	
	# note refs.
	if tmp[0] not in refs:
		refs[tmp[0]] = idx
		idx += 1
	
	# add to array.
	queries[j]['type'] = False
	queries[j]['ref'] = refs[tmp[0]]
	queries[j]['start'] = int(tmp[3])
	queries[j]['stop'] = int(tmp[4])
	queries[j]['info'] = j
	j += 1
	
	#if i >= debug_sz: break


# Create overlap object.
num_ref = len(refs)
olap = overlap.Overlap(targets, queries, num_ref, verbose=True)

# Find weak overlap.
weako = olap.weak_overlap(min_ov=prec_ov)

# Mark query with overlapping targets.
if just_mark != False:
	# We just print the overlapping queries and targets.
	print "printing marked query."
	fout = open(just_mark, "wb")
	
	for pair in weako:
		# get entries.
		tline = target_lines[targets[pair[0]]['info']]
		qline = query_lines[queries[pair[1]]['info']]
		
		# extend 9.
		if qline[8] != "":
			qline[8] += " ; "
		qline[8] += tline[8]
		
		# write to file.
		fout.write('\t'.join(qline) + '\n')
	
	fout.close()

# do if hits or misses.
if hits == True or misses == True:	
	# Flag overlapping edges.
	print "flagging overlaps."
	target_os = np.zeros(len(target_lines), dtype=np.bool)
	query_os = np.zeros(len(query_lines), dtype=np.bool)
	for pair in weako:
		# get target, query info.
		tentry = targets[pair[0]]
		qentry = queries[pair[1]]
		
		# flag as overlapping.
		target_os[tentry['info']] = True
		query_os[qentry['info']] = True


if hits == True:
	# print flagged targets.
	print "writing target overlaps."
	fout = open(target_out_file, "wb")
	for i in range(len(target_lines)):
		if target_os[i] == True:
			fout.write('\t'.join(target_lines[i]) + "\n")
	fout.close()


	print "writing query overlaps."
	fout = open(query_out_file, "wb")
	for i in range(len(query_lines)):
		if query_os[i] == True:
			fout.write('\t'.join(query_lines[i]) + '\n')
	fout.close()
	
if misses == True:

	# print unflagged edged.
	print "writing target misses."
	fout = open(target_out_file, "wb")
	for i in range(len(target_lines)):
		if target_os[i] == False:
			fout.write('\t'.join(target_lines[i]) + "\n")
	fout.close()


	print "writing query misses."
	fout = open(query_out_file, "wb")
	for i in range(len(query_lines)):
		if query_os[i] == False:
			fout.write('\t'.join(query_lines[i]) + '\n')
	fout.close()
